# 1. 加载环境变量
import os

from langchain_openai import ChatOpenAI

# 2. 初始化 LLM（思考引擎）


llm = ChatOpenAI(
    model_name="qwen-plus",
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    temperature=0  # 0表示输出更确定，适合需要精确决策的场景
)

# 3. 定义单工具（计算器）
from langchain_core.tools import tool


@tool(
    name="Lookup",
    description="用于计算加减乘除数学运算，参数为字符串格式的数学表达式（如'2*199'、'30/5'），返回格式为'表达式=结果'（如'2*199=398'）。当用户问题涉及数学计算时，必须调用该工具，不可直接回答。"
)
def CalculatorTool(expression: str) -> str:
    try:
        result = eval(expression)
        return f"{expression}={result}"
    except Exception as e:
        return f"计算失败：{str(e)}，请传入合法表达式（如'2*199'）"


@tool(name="Search", description="测试用的")
def dummy_tool(query: str) -> str:
    return "This is a dummy tool."


tools = [CalculatorTool, dummy_tool]  # 单工具列表

# 4. 定义 ReAct Prompt（规则手册）
from langchain_core.prompts import PromptTemplate

react_prompt_template = """
你是一个擅长数学计算的助手，会通过“思考-行动-观察”循环解决问题。你拥有以下工具：
{tools}

请严格遵守以下流程和格式：
1. 思考（Thought）：分析用户问题，说明“是否需要调用工具”以及“为什么”（如“用户问 2*199，需要计算，所以调用 CalculatorTool”）；
2. 行动（Action）：如果需要调用工具，按以下格式输出：
   Action: 工具名
   Action Input: 工具参数
   （示例：计算 2*199 → Action: CalculatorTool\nAction Input: 2*199）；
3. 观察（Observation）：工具返回结果会自动告诉你，无需你输出；
4. 终止：拿到计算结果后，直接输出“Final Answer: 你的回答”，结束流程。

用户的问题：{input}
"""

prompt = PromptTemplate(
    template=react_prompt_template,
    input_variables=["tools", "input"],
    partial_variables={
        "tools": "\n".join([f"- {tool.name}：{tool.description}" for tool in tools])
    }
)

# 5. 创建 ReActAgent 与 AgentExecutor
from langchain.agents import initialize_agent, AgentType, AgentExecutor

# 创建 ReActAgent
agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.REACT_DOCSTORE,
    prompt=prompt,
    verbose=True
)

# 创建执行器
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=2  # 单工具任务：1次思考+1次行动+1次观察，2次迭代足够
)

# 6. 测试 ReActAgent（计算 2*199 + 5）
if __name__ == "__main__":
    user_input = "计算 2*199 + 5 的结果是多少？"
    result = agent_executor.invoke({"input": user_input})
    # 输出最终答案
    print("\n===== 最终结果 =====")
    print(result["output"])
